package first.group.ssh.packet;

import src.AESCipherCTR;

import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import javax.crypto.Cipher;
import javax.crypto.Mac;
import javax.crypto.ShortBufferException;

public class PacketSerializer {

    private AESCipherCTR cipher = null;
    private Mac mac = null;
    private int seq = 0;

    private void serializeUnencrypted(Packet packet, ByteBuffer buffer) {
        var startPos = buffer.position();

        buffer.position(buffer.position() + 6);
        var dataStart = buffer.position();
        packet.write(buffer);
        var dataEnd = buffer.position();
        var dataLen = dataEnd - dataStart;
        
        var blockSize = 8;
        if (cipher != null) {
            blockSize = cipher.getBlockSize();
        }

        var packetLen = 2 + dataLen;
        var rawLen = 4 + packetLen;
        int totalLen, paddingLen;
        if (rawLen % blockSize < 4) {
            paddingLen = blockSize - (rawLen % blockSize);
        } else {
            paddingLen = blockSize * 2 - (rawLen % blockSize);
        }
        totalLen = rawLen + paddingLen;
        assert (totalLen % blockSize == 0);
        assert (paddingLen >= 4);

        buffer.position(startPos);
        buffer.putInt(totalLen - 4);
        buffer.put((byte) paddingLen);
        buffer.put((byte) packet.getType());

        buffer.position(dataEnd);
        for (int i = 0; i < paddingLen; i++) {
            buffer.put((byte) 0);
        }
    }

    public void setCipher(AESCipherCTR cipher, Mac mac) {
        this.cipher = cipher;
        this.mac = mac;
    }

    private void serializeEncrypted(Packet packet, ByteBuffer buffer) {
        var clearTextBuffer = ByteBuffer.allocate(65536);
        clearTextBuffer.order(ByteOrder.BIG_ENDIAN);
        serializeUnencrypted(packet, clearTextBuffer);
        clearTextBuffer.flip();
        mac.reset();
        
        var seqBuffer = ByteBuffer.allocate(4);
        seqBuffer.putInt(seq);
        seqBuffer.flip();
        mac.update(seqBuffer);
        mac.update(clearTextBuffer);

        byte[] macData = mac.doFinal();
        clearTextBuffer.flip();

        cipher.update(clearTextBuffer, buffer);
        buffer.put(macData);
    }

    public ByteBuffer serialize(Packet packet) {
        var buffer = ByteBuffer.allocate(65536);
        buffer.order(ByteOrder.BIG_ENDIAN);
        if (cipher != null) {
            serializeEncrypted(packet, buffer);
        } else {
            serializeUnencrypted(packet, buffer);
        }
        buffer.flip();
        seq++;
        return buffer;
    }
}
